import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms

# 数据预处理
transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,))])

# 加载MNIST数据集
train_dataset = datasets.MNIST(root='./data', train=True, download=True, transform=transform)
test_dataset = datasets.MNIST(root='./data', train=False, download=True, transform=transform)

train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=64, shuffle=True)
test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=1000, shuffle=False)

# 定义简单的神经网络模型
class SimpleNN(nn.Module):
    def __init__(self):
        super(SimpleNN, self).__init__()
        self.flatten = nn.Flatten()
        self.fc1 = nn.Linear(28 * 28, 128)
        self.relu = nn.ReLU()
        self.dropout = nn.Dropout(0.2)
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        x = self.flatten(x)
        x = self.fc1(x)
        x = self.relu(x)
        x = self.dropout(x)
        x = self.fc2(x)
        return x

model = SimpleNN()

# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

# 训练模型
for epoch in range(5):  # 训练5个epoch
    model.train()
    running_loss = 0.0
    for images, labels in train_loader:
        optimizer.zero_grad()
        outputs = model(images)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        running_loss += loss.item()
    print(f'Epoch [{epoch + 1}/5], Loss: {running_loss / len(train_loader)}')

# 评估模型
model.eval()
correct = 0
total = 0
with torch.no_grad():
    for images, labels in test_loader:
        outputs = model(images)
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

print(f'Accuracy of the network on the 10000 test images: {100 * correct / total}%')

# pip install torch torchvision
# python mnist_pytorch.py
# 1.数据预处理：使用torchvision.transforms进行数据归一化处理。
# 2.加载MNIST数据集：使用torchvision.datasets加载数据集，并使用DataLoader创建数据加载器。
# 3.定义模型：创建一个继承自nn.Module的简单神经网络类SimpleNN。
# 4.定义损失函数和优化器：使用交叉熵损失函数nn.CrossEntropyLoss和Adam优化器optim.Adam。
# 5.训练模型：在每个epoch中，通过前向传播计算输出，计算损失，反向传播更新权重。
# 6.评估模型：在测试集上评估模型性能，计算准确率。

# 官方网站：http://yann.lecun.com/exdb/mnist/。
# 这是 MNIST 数据集的原始来源，提供四个 gz 压缩文件，分别是训练图像数据、训练标签数据、测试图像数据和测试标签数据。
# https://gitcode.com/open-source-toolkit/54201
# https://gitcode.com/open-source-toolkit/39c42
# https://gitcode.com/open-source-toolkit/4b549
# 官网提供下载，但由于是国外的服务器，下载速度很慢。这里提供百度网盘下载地址：
# 链接：https://pan.baidu.com/s/17KUWe8JdQBHsAg3B4m5SNg 
# 提取码：wyxn

# torchvision是做什么用的？
# torchvision 是 PyTorch 生态中的核心组件，主要用于处理图像相关的任务，以下是其具体用途介绍：
# 数据加载与预处理：内置了多种常用的图像数据集，如 MNIST、CIFAR-10、ImageNet 等，方便用户快速获取和加载数据。同时，提供了丰富的图像预处理操作，如裁剪、缩放、归一化、旋转、翻转等，这些操作可以单独使用或组合使用，形成复杂的预处理流程，以满足不同模型对输入数据的要求。
# 模型构建：包含了许多经典的预训练模型，如 AlexNet、VGG、ResNet、Inception 等。这些模型在大规模图像数据上进行了预训练，具有较高的准确性和泛化能力，用户可以将其作为基础模型进行微调或迁移学习，快速构建自己的图像识别、分类等模型。
# 模型训练与优化：提供了一些针对模型训练的优化策略和工具，如优化器、学习率调度器等，帮助用户更好地调整模型的训练过程，提高模型的性能和收敛速度。
# 综上所述，torchvision 为计算机视觉任务提供了全面而强大的支持，极大地简化了从数据处理到模型构建、训练及优化的整个流程。

# 执行过程记录
# E:\train\mnist>python mnist_pytorch.py
# Downloading http://yann.lecun.com/exdb/mnist/train-images-idx3-ubyte.gz
# Failed to download (trying next):
# <urlopen error [WinError 10060] 由于连接方在一段时间后没有正确答复或连接的主机没有反应，连接尝试失败。>
#
# Downloading https://ossci-datasets.s3.amazonaws.com/mnist/train-images-idx3-ubyte.gz
# Downloading https://ossci-datasets.s3.amazonaws.com/mnist/train-images-idx3-ubyte.gz to ./data\MNIST\raw\train-images-idx3-ubyte.gz
# 100%|██████████████████████████████████████████████████████████████████████████████| 9.91M/9.91M [00:39<00:00, 251kB/s]
# Extracting ./data\MNIST\raw\train-images-idx3-ubyte.gz to ./data\MNIST\raw
#
# Downloading http://yann.lecun.com/exdb/mnist/train-labels-idx1-ubyte.gz
# Failed to download (trying next):
# <urlopen error [WinError 10060] 由于连接方在一段时间后没有正确答复或连接的主机没有反应，连接尝试失败。>
#
# Downloading https://ossci-datasets.s3.amazonaws.com/mnist/train-labels-idx1-ubyte.gz
# Downloading https://ossci-datasets.s3.amazonaws.com/mnist/train-labels-idx1-ubyte.gz to ./data\MNIST\raw\train-labels-idx1-ubyte.gz
# 100%|█████████████████████████████████████████████████████████████████████████████| 28.9k/28.9k [00:00<00:00, 95.0kB/s]
# Extracting ./data\MNIST\raw\train-labels-idx1-ubyte.gz to ./data\MNIST\raw
#
# Downloading http://yann.lecun.com/exdb/mnist/t10k-images-idx3-ubyte.gz
# Failed to download (trying next):
# <urlopen error [WinError 10060] 由于连接方在一段时间后没有正确答复或连接的主机没有反应，连接尝试失败。>
#
# Downloading https://ossci-datasets.s3.amazonaws.com/mnist/t10k-images-idx3-ubyte.gz
# Downloading https://ossci-datasets.s3.amazonaws.com/mnist/t10k-images-idx3-ubyte.gz to ./data\MNIST\raw\t10k-images-idx3-ubyte.gz
# 100%|█████████████████████████████████████████████████████████████████████████████| 1.65M/1.65M [00:23<00:00, 69.6kB/s]
# Extracting ./data\MNIST\raw\t10k-images-idx3-ubyte.gz to ./data\MNIST\raw

# Downloading http://yann.lecun.com/exdb/mnist/t10k-labels-idx1-ubyte.gz
# Failed to download (trying next):
# <urlopen error [WinError 10060] 由于连接方在一段时间后没有正确答复或连接的主机没有反应，连接尝试失败。>
# Downloading https://ossci-datasets.s3.amazonaws.com/mnist/t10k-labels-idx1-ubyte.gz
# Downloading https://ossci-datasets.s3.amazonaws.com/mnist/t10k-labels-idx1-ubyte.gz to ./data\MNIST\raw\t10k-labels-idx1
# -ubyte.gz
# 100%|█████████████████████████████████████████████████████████████████████████████| 4.54k/4.54k [00:00<00:00, 1.44MB/s]
# Extracting ./data\MNIST\raw\t10k-labels-idx1-ubyte.gz to ./data\MNIST\raw
# Epoch [1/5], Loss: 0.41979895378830334
# Epoch [2/5], Loss: 0.22694163110210444
# Epoch [3/5], Loss: 0.18395106531361907
# Epoch [4/5], Loss: 0.15901574041984323
# Epoch [5/5], Loss: 0.1413092447081402
# Accuracy of the network on the 10000 test images: 96.43%
# E:\train\mnist>python mnist_gpu.py
# Using device: cuda
# Epoch [1/5], Loss: 0.4024941038602451
# Epoch [2/5], Loss: 0.21319795640181505
# Epoch [3/5], Loss: 0.17209210851068882
# Epoch [4/5], Loss: 0.14925990098996011
# Epoch [5/5], Loss: 0.13816061745217043
# Accuracy of the network on the 10000 test images: 96.68%
# E:\train\mnist>

# 训练的环境,可测试使用的是已安装的，未测试版本兼容
# python 3.9.0
# pytorch版本：2.5.1+cu124